{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 深度概率编程\n",
    "\n",
    "## 概述\n",
    "\n",
    "深度学习模型具有强大的拟合能力，而贝叶斯理论具有很好的可解释能力。MindSpore深度概率编程（MindSpore Deep Probabilistic Programming, MDP）将深度学习和贝叶斯学习结合，通过设置网络权重为分布、引入隐空间分布等，可以对分布进行采样前向传播，由此引入了不确定性，从而增强了模型的鲁棒性和可解释性。MDP不仅包含通用、专业的概率学习编程语言，适用于“专业”用户，而且支持使用开发深度学习模型的逻辑进行概率编程，让初学者轻松上手；此外，还提供深度概率学习的工具箱，拓展贝叶斯应用功能。\n",
    "\n",
    "本章将详细介绍深度概率编程在MindSpore上的应用。在动手进行实践之前，确保，你已经正确安装了MindSpore 0.7.0-beta及其以上版本。\n",
    "\n",
    "> 本例适用于GPU和Ascend环境。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 环境准备\n",
    "\n",
    "设置训练模式为图模式，计算平台为GPU。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mindspore import context\n",
    "\n",
    "context.set_context(mode=context.GRAPH_MODE, save_graphs=False, device_target=\"GPU\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 数据准备\n",
    "\n",
    "### 下载数据集\n",
    "下载MNIST数据集并解压到指定位置，执行如下命令："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2020-12-17 15:38:35--  https://obs.dualstack.cn-north-4.myhuaweicloud.com/mindspore-website/notebook/datasets/MNIST_Data.zip\n",
      "Resolving proxy-notebook.modelarts-dev-proxy.com (proxy-notebook.modelarts-dev-proxy.com)... 192.168.0.172\n",
      "Connecting to proxy-notebook.modelarts-dev-proxy.com (proxy-notebook.modelarts-dev-proxy.com)|192.168.0.172|:8083... connected.\n",
      "Proxy request sent, awaiting response... 200 OK\n",
      "Length: 10754903 (10M) [application/zip]\n",
      "Saving to: ‘MNIST_Data.zip’\n",
      "\n",
      "MNIST_Data.zip      100%[===================>]  10.26M  48.5MB/s    in 0.2s    \n",
      "\n",
      "2020-12-17 15:38:35 (48.5 MB/s) - ‘MNIST_Data.zip’ saved [10754903/10754903]\n",
      "\n",
      "Archive:  MNIST_Data.zip\n",
      "   creating: ./datasets/MNIST_Data/test/\n",
      "  inflating: ./datasets/MNIST_Data/test/t10k-images-idx3-ubyte  \n",
      "  inflating: ./datasets/MNIST_Data/test/t10k-labels-idx1-ubyte  \n",
      "   creating: ./datasets/MNIST_Data/train/\n",
      "  inflating: ./datasets/MNIST_Data/train/train-images-idx3-ubyte  \n",
      "  inflating: ./datasets/MNIST_Data/train/train-labels-idx1-ubyte  \n",
      "./datasets/MNIST_Data/\n",
      "├── test\n",
      "│   ├── t10k-images-idx3-ubyte\n",
      "│   └── t10k-labels-idx1-ubyte\n",
      "└── train\n",
      "    ├── train-images-idx3-ubyte\n",
      "    └── train-labels-idx1-ubyte\n",
      "\n",
      "2 directories, 4 files\n"
     ]
    }
   ],
   "source": [
    "!wget -N https://obs.dualstack.cn-north-4.myhuaweicloud.com/mindspore-website/notebook/datasets/MNIST_Data.zip\n",
    "!unzip -o MNIST_Data.zip -d ./datasets/\n",
    "!tree ./datasets/MNIST_Data/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 定义数据集增强方法\n",
    "\n",
    "MNIST数据集的原始训练数据集是60000张$28\\times28$像素的单通道数字图片，本次训练用到的含贝叶斯层的LeNet5网络接收到训练数据的张量为`(32,1,32,32)`，通过自定义create_dataset函数将原始数据集增强为适应训练要求的数据，具体的增强操作解释可参考官网快速入门[实现一个图片分类应用](https://www.mindspore.cn/tutorial/training/zh-CN/master/quick_start/quick_start.html)。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import mindspore.dataset.vision.c_transforms as CV\n",
    "import mindspore.dataset.transforms.c_transforms as C\n",
    "from mindspore.dataset.vision import Inter\n",
    "from mindspore import dataset as ds\n",
    "\n",
    "def create_dataset(data_path, batch_size=32, repeat_size=1,\n",
    "                   num_parallel_workers=1):\n",
    "    # define dataset\n",
    "    mnist_ds = ds.MnistDataset(data_path)\n",
    "\n",
    "    # define some parameters needed for data enhancement and rough justification\n",
    "    resize_height, resize_width = 32, 32\n",
    "    rescale = 1.0 / 255.0\n",
    "    shift = 0.0\n",
    "    rescale_nml = 1 / 0.3081\n",
    "    shift_nml = -1 * 0.1307 / 0.3081\n",
    "\n",
    "    # according to the parameters, generate the corresponding data enhancement method\n",
    "    c_trans = [\n",
    "        CV.Resize((resize_height, resize_width), interpolation=Inter.LINEAR),\n",
    "        CV.Rescale(rescale_nml, shift_nml),\n",
    "        CV.Rescale(rescale, shift),\n",
    "        CV.HWC2CHW()\n",
    "    ]\n",
    "    type_cast_op = C.TypeCast(mstype.int32)\n",
    "\n",
    "    # using map to apply operations to a dataset\n",
    "    mnist_ds = mnist_ds.map(operations=type_cast_op, input_columns=\"label\", num_parallel_workers=num_parallel_workers)\n",
    "    mnist_ds = mnist_ds.map(operations=c_trans, input_columns=\"image\", num_parallel_workers=num_parallel_workers)\n",
    "\n",
    "    \n",
    "    # process the generated dataset\n",
    "    buffer_size = 10000\n",
    "    mnist_ds = mnist_ds.shuffle(buffer_size=buffer_size)\n",
    "    mnist_ds = mnist_ds.batch(batch_size, drop_remainder=True)\n",
    "    mnist_ds = mnist_ds.repeat(repeat_size)\n",
    "\n",
    "    return mnist_ds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 定义深度神经网络\n",
    "\n",
    "在经典LeNet5网络中，数据经过如下计算过程：卷积1->激活->池化->卷积2->激活->池化->降维->全连接1->全连接2->全连接3。  \n",
    "本例中将引入概率编程方法，将卷积1和全连接1两个计算层改造成贝叶斯层，构造成含贝叶斯层的LeNet5网络。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mindspore.common.initializer import Normal\n",
    "import mindspore.nn as nn\n",
    "from mindspore.nn.probability import bnn_layers\n",
    "import mindspore.ops as ops\n",
    "from mindspore import dtype as mstype\n",
    "\n",
    "\n",
    "class BNNLeNet5(nn.Cell):\n",
    "    def __init__(self, num_class=10):\n",
    "        super(BNNLeNet5, self).__init__()\n",
    "        self.num_class = num_class\n",
    "        self.conv1 = bnn_layers.ConvReparam(1, 6, 5, stride=1, padding=0, has_bias=False, pad_mode=\"valid\")\n",
    "        self.conv2 = nn.Conv2d(6, 16, 5, pad_mode='valid')\n",
    "        self.fc1 = bnn_layers.DenseReparam(16 * 5 * 5, 120)\n",
    "        self.fc2 = nn.Dense(120, 84, weight_init=Normal(0.02))\n",
    "        self.fc3 = nn.Dense(84, self.num_class)\n",
    "        self.relu = nn.ReLU()\n",
    "        self.max_pool2d = nn.MaxPool2d(kernel_size=2, stride=2)\n",
    "        self.flatten = nn.Flatten()\n",
    "\n",
    "    def construct(self, x):\n",
    "        x = self.max_pool2d(self.relu(self.conv1(x)))\n",
    "        x = self.max_pool2d(self.relu(self.conv2(x)))\n",
    "        x = self.flatten(x)\n",
    "        x = self.relu(self.fc1(x))\n",
    "        x = self.relu(self.fc2(x))\n",
    "        x = self.fc3(x) \n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "本例中将卷积层1和全连接1两个计算层换成了贝叶斯卷积层`bnn_layers.ConvReparam`和贝叶斯全连接层`bnn_layers.DenseReparam`。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 定义训练网络\n",
    "\n",
    "定义训练网络并进行训练。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 1 step: 1, loss is 2.3022718\n",
      "epoch: 1 step: 2, loss is 2.3022223\n",
      "epoch: 1 step: 3, loss is 2.3028727\n",
      "epoch: 1 step: 4, loss is 2.3034232\n",
      "epoch: 1 step: 5, loss is 2.3019493\n",
      "epoch: 1 step: 6, loss is 2.3017588\n",
      "... ...\n",
      "epoch: 1 step: 1866, loss is 0.097549394\n",
      "epoch: 1 step: 1867, loss is 0.082386635\n",
      "epoch: 1 step: 1868, loss is 0.027000971\n",
      "epoch: 1 step: 1869, loss is 0.026424333\n",
      "epoch: 1 step: 1870, loss is 0.19351783\n",
      "epoch: 1 step: 1871, loss is 0.02400064\n",
      "epoch: 1 step: 1872, loss is 0.3389563\n",
      "epoch: 1 step: 1873, loss is 0.004886848\n",
      "epoch: 1 step: 1874, loss is 0.020785151\n",
      "epoch: 1 step: 1875, loss is 0.33145565\n"
     ]
    }
   ],
   "source": [
    "from mindspore.nn import TrainOneStepCell\n",
    "from mindspore import Tensor, Model\n",
    "from mindspore.train.callback import ModelCheckpoint, CheckpointConfig, LossMonitor\n",
    "from mindspore.nn.metrics import Accuracy\n",
    "from mindspore.nn.loss import SoftmaxCrossEntropyWithLogits\n",
    "import os\n",
    "\n",
    "\n",
    "lr = 0.01\n",
    "momentum = 0.9\n",
    "model_path = \"./models/ckpt/probability_bnnlenet5/\"\n",
    "\n",
    "# clean old run files\n",
    "os.system(\"rm -f {0}*.meta {0}*.ckpt\".format(model_path))\n",
    "network = BNNLeNet5()\n",
    "criterion = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction=\"mean\")\n",
    "optimizer = nn.Momentum(network.trainable_params(), lr, momentum)\n",
    "model = Model(network, criterion, optimizer, metrics={\"Accuracy\": Accuracy()} )\n",
    "\n",
    "config_ck = CheckpointConfig(save_checkpoint_steps=1875, keep_checkpoint_max=16)\n",
    "ckpoint_cb = ModelCheckpoint(prefix=\"checkpoint_lenet\", directory=model_path, config=config_ck)\n",
    "\n",
    "ds_train_path = \"./datasets/MNIST_Data/train/\"\n",
    "train_set = create_dataset(ds_train_path, 32, 1)\n",
    "model.train(1, train_set, callbacks=[ckpoint_cb, LossMonitor()])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "训练完成后会在对应的路径上生成`.ckpt`为后缀的权重参数文件和`.meta`为后缀的计算图文件。  \n",
    "其路径结构为："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "./models/ckpt/probability_bnnlenet5/\n",
      "├── checkpoint_lenet-1_1875.ckpt\n",
      "└── checkpoint_lenet-graph.meta\n",
      "\n",
      "0 directories, 2 files\n"
     ]
    }
   ],
   "source": [
    "!tree $model_path"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 验证模型精度\n",
    "\n",
    "载入验证数据集，并验证含有贝叶斯层的LeNet5网络模型的精度。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'Accuracy': 0.9730568910256411}\n"
     ]
    }
   ],
   "source": [
    "ds_eval_path = \"./datasets/MNIST_Data/test/\"\n",
    "test_set = create_dataset(ds_eval_path, 32, 1)\n",
    "acc = model.eval(test_set)\n",
    "print(acc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "模型精度大于0.95，证明模型效果良好。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 总结\n",
    "\n",
    "本例使用了深度概率编程在经典LeNet5深度神经网络中应用，含有贝叶斯层的LeNet5网络和原本的LeNet5网络的训练体验过程极其相似，有心的用户可以对比两者在训练收敛效率，稳定性等方面的不同，是否体现了概述中深度概率编程的优点。  \n",
    "当然深度概率编程近年来最激动人心的是在CVAE以及GAN等生成网络中的应用，这使我们在拥有了以假乱真的数据生成能力，接下来一篇就以CVAE网络体验介绍深度概率编程的另一种应用。"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "MindSpore-1.0.1",
   "language": "python",
   "name": "mindspore-1.0.1"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
